{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 06. GRIN medium and analysing refraction\n",
    "\n",
    "submitted by [substancia](https://github.com/substancia), adapted by [flaport](https://github.com/flaport)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import fdtd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid = fdtd.Grid(shape=(9.3e-6, 15.5e-6, 1), grid_spacing=77.5e-9)\n",
    "# x boundaries\n",
    "grid[0:10, :, :] = fdtd.PML(name=\"pml_xlow\")\n",
    "grid[-10:, :, :] = fdtd.PML(name=\"pml_xhigh\")\n",
    "# y boundaries\n",
    "grid[:, 0:10, :] = fdtd.PML(name=\"pml_ylow\")\n",
    "grid[:, -10:, :] = fdtd.PML(name=\"pml_yhigh\")\n",
    "simfolder = grid.save_simulation(\"GRIN\")  # initializing environment to save simulation data\n",
    "print(simfolder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Objects\n",
    "defining a graded refractive index slab, with homogenous slab extensions outwards from both ends"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n0, theta, t = 1, 30, 0.5\n",
    "for i in range(50):\n",
    "    x = i * 0.08\n",
    "    epsilon = n0 + x * np.sin(np.radians(theta)) / t\n",
    "    epsilon = epsilon ** 0.5\n",
    "    grid[\n",
    "        5.1e-6:5.6e-6, (5 + i * 0.08) * 1e-6 : (5.08 + i * 0.08) * 1e-6, 0\n",
    "    ] = fdtd.Object(permittivity=epsilon, name=\"object\" + str(i))\n",
    "\n",
    "# homogenous slab extensions\n",
    "grid[5.1e-6:5.6e-6, 0.775e-6:5e-6, 0] = fdtd.Object(\n",
    "    permittivity=n0 ** 2, name=\"objectLeft\"\n",
    ")\n",
    "grid[5.1e-6:5.6e-6, 9e-6 : (15.5 - 0.775) * 1e-6, 0] = fdtd.Object(\n",
    "    permittivity=epsilon, name=\"objectRight\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Source\n",
    "using a pulse (hanning window pulse)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid[3.1e-6, 1.5e-6:14e-6, 0] = fdtd.LineSource(period=1550e-9 / (3e8), name=\"source\", pulse=True, cycle=3, hanning_dt=4e-15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Detectors\n",
    "using a linear array of LineDetector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(-4, 8):\n",
    "    grid[5.8e-6, 84 + 4 * i : 86 + 4 * i, 0] = fdtd.LineDetector(name=\"detector\" + str(i))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Saving grid geometry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(os.path.join(\"./fdtd_output\", grid.folder, \"grid.txt\"), \"w\") as f:\n",
    "    f.write(str(grid))\n",
    "    wavelength = 3e8/grid.source.frequency\n",
    "    wavelengthUnits = wavelength/grid.grid_spacing\n",
    "    GD = np.array([grid.x, grid.y, grid.z])\n",
    "    gridRange = [np.arange(x/grid.grid_spacing) for x in GD]\n",
    "    objectRange = np.array([[gridRange[0][x.x], gridRange[1][x.y], gridRange[2][x.z]] for x in grid.objects], dtype=object).T\n",
    "    f.write(\"\\n\\nGrid details (in wavelength scale):\")\n",
    "    f.write(\"\\n\\tGrid dimensions: \")\n",
    "    f.write(str(GD/wavelength))\n",
    "    f.write(\"\\n\\tSource dimensions: \")\n",
    "    f.write(str(np.array([grid.source.x[-1] - grid.source.x[0] + 1, grid.source.y[-1] - grid.source.y[0] + 1, grid.source.z[-1] - grid.source.z[0] + 1])/wavelengthUnits))\n",
    "    f.write(\"\\n\\tObject dimensions: \")\n",
    "    f.write(str([(max(map(max, x)) - min(map(min, x)) + 1)/wavelengthUnits for x in objectRange]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import clear_output  # only necessary in jupyter notebooks\n",
    "\n",
    "for i in range(100):\n",
    "    grid.step()  # running simulation 1 timestep a time and animating\n",
    "    if i % 5 == 0:\n",
    "        # saving frames during visualization\n",
    "        grid.visualize(z=0, animate=True, index=i, save=True, folder=simfolder)\n",
    "        plt.title(f\"{i:3.0f}\")\n",
    "        clear_output(wait=True)  # only necessary in jupyter notebooks\n",
    "grid.save_data()  # saving detector readings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can generate a video with ffmpeg:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    video_path = grid.generate_video(delete_frames=False)  # rendering video from saved frames\n",
    "except:\n",
    "    video_path = \"\"\n",
    "    print(\"ffmpeg not installed?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if video_path:\n",
    "    from IPython.display import Video\n",
    "    display(Video(video_path, embed=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyse\n",
    "analysing data stored by above simulation to find intensity profile and time-of-arrival plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dic = np.load(os.path.join(simfolder, \"detector_readings.npz\"))\n",
    "import warnings; warnings.filterwarnings(\"ignore\") # TODO: fix plot_detection to prevent warnings\n",
    "fdtd.plot_detection(dic)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
